import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter


# 测试手写的input
def t_input():
    # 定义一个5x5的输入矩阵（模拟单通道图像）
    input = torch.tensor([[1, 2, 0, 3, 1],
                          [0, 1, 2, 3, 1],
                          [1, 2, 1, 0, 0],
                          [5, 2, 3, 1, 1],
                          [2, 1, 0, 1, 1]
                          ])

    """
    将输入张量 input 重新调整为 4 维形状 (batch_size, channels, height, width)，这是 PyTorch 卷积/池化层要求的标准输入格式。
    -1	自动计算的批大小（batch_size）	系统根据元素总数推算   计算公式：总元素数 / (1 * 5 * 5)
    """
    input = torch.reshape(input, (-1, 1, 5, 5))
    print(f"输入input：{input.shape}\n")

    output = net(input)
    print(f"输出output：{output}\n")
    print(f"输出output形状：{output.shape}\n")


# 测试torchvision中的dataset
def t_dataset():
    dataset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True,
                                           transform=torchvision.transforms.ToTensor())

    train_loader = torch.utils.data.DataLoader(dataset, batch_size=64)

    writer = SummaryWriter(log_dir="logs/log_maxpool")
    step = 0
    for data in train_loader:
        inputs, labels = data
        writer.add_images(tag="inputs", img_tensor=inputs, global_step=step)

        outputs = net(inputs)
        writer.add_images(tag="outputs", img_tensor=outputs, global_step=step)

        step = step + 1

    writer.close()


if __name__ == '__main__':
    # 网络定义
    class Net(torch.nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.maxpool1 = torch.nn.MaxPool2d(kernel_size=3, ceil_mode=True)

        def forward(self, x):
            x = self.maxpool1(x)
            return x


    """
    net 是在 if __name__ == '__main__': 代码块中定义的，但它是一个全局变量（因为 Python 的模块作用域规则）。
    因此，t_input() 和 t_dataset() 可以直接访问 net，无需显式传递。
    """
    net = Net()
    # print(f"网络net：{net}\n")

    # t_input()

    t_dataset()
